import torch
from torch import nn
from torch.nn import functional as F
from torch import optim
import torchvision
from matplotlib import pyplot as plt
from utils import plot_image,plot_curve ,one_hot

batch_size = 512

train_loader = torch.utils.data.DataLoader(
    torchvision.datasets.MNIST('mnist_data',train = True,download=True,
                               transform=torchvision.transforms.Compose(
                                   [
                                       torchvision.transforms.ToTensor(),
                                       torchvision.transforms.Normalize(
                                           (0.1307,),(0.3081,)
                                       )
                                   ]
                               )),
    batch_size=batch_size,shuffle=True
)
test_loader = torch.utils.data.DataLoader(
    torchvision.datasets.MNIST('mnist_data',train = True,download=True,
                               transform=torchvision.transforms.Compose(
                                   [
                                       torchvision.transforms.ToTensor(),
                                       torchvision.transforms.Normalize(
                                           (0.1307,),(0.3081,)
                                       )
                                   ]
                               )),
    batch_size=batch_size,shuffle=True
)
x,y = next(iter(train_loader))
plot_image(x,y,'image sample')

class Net(nn.Module):
    def __init__(self):
        super(Net,self).__init__()

        self.fc1 = nn.Linear(28*28,256)
        self.fc2 = nn.Linear (256,64)
        self.fc3 = nn.Linear(64,10)

    def forward(self,x):
        x=F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = self.fc3(x)

        return x

net = Net()
optimizer = optim.SGD(net.parameters(),lr=0.01,momentum=0.9)
for epoch in range(3):
    for batch_idx,(x,y) in enumerate(train_loader):
        print(x.size())
        x = x.view(x.size(0),1*28*28)
        print(x.size())
        out = net(x)
        y_onehot = one_hot(y)
        loss = F.mse_loss(out,y_onehot)
        loss.backward()
        optimizer.step()

        if batch_idx % 10 == 0 :
            print(epoch,batch_idx,loss.item())




